'''
Created on 9. 1, 2021
@author: Kailun Yang (yangkailunysu@163.com)
'''
__author__ = "Kailun Yang"

import random
import numpy as np
import xlearn as xl
import pandas as pd
from collections import defaultdict
import heapq
import os
from utils.parser import parse_args
from sklearn import preprocessing

if __name__ == '__main__':
    """fix the random seed"""
    seed = 2020
    random.seed(seed)
    np.random.seed(seed)
    args = parse_args()

    spu_off = 200000
    period_off = 395244
    aor_off = 395249
    spu_num = 195244
    poi_off = 395260
    avg_pay_amt_off = 425184
    spu_price_qujian_off = 425189

    id = args.model + '_' + str(args.run_id)
    bprmf_out = '3_top_50_submit_KGIN_data_KGIN_RI_test'
    mode = 'train'
    train_file = 'submit-train_5'

    data = pd.read_table('../../data/three_week_data.txt', sep='\t')
    test = pd.read_table('../../data/orders_test_spu.txt', sep='\t')
    train = data

    # 根据用户的购买记录，生成每个spu所属的poi的字典
    spu_of_poi_dict = defaultdict(list)
    for poi, spu in zip(train['wm_poi_id'], train['wm_food_spu_id']):
        spu_of_poi_dict[spu].append(poi)
    for spu in spu_of_poi_dict:
        if len(spu_of_poi_dict[spu]) > 1:
            spu_of_poi_dict[spu] = list(set(spu_of_poi_dict[spu]))

    # 读入用户在训练集中点过的菜品，作为后面抽样负样本的根据
    user_spu = defaultdict(list)
    for line in open('../solution_1/data/submit_KGIN_data/train.txt'):
        line = line.strip()
        lines = line.split(' ')
        lines = [int(a) for a in lines]
        user_spu[lines[0]].extend(list(set(lines[1:])))

    users = pd.read_table('../../data/users.txt', sep='\t')
    # 对用户的历史单价做labelEncoder后建立字典
    avg_pay_amt_label = preprocessing.LabelEncoder().fit_transform(users.loc[:, 'avg_pay_amt'])
    user_avg_pay_amt_dict = {}
    for user, avg_pay_amt in zip(range(0, 200000), avg_pay_amt_label):
        user_avg_pay_amt_dict[user] = avg_pay_amt

    spus = pd.read_table('../../data/spus.txt', sep='\t')

    def zero():
        return 0

    # 构造菜品的价格区间这个特征
    spu_price_qujian = defaultdict(zero)
    for spu, price in zip(spus['wm_food_spu_id'], spus['price']):
        if '.' in str(price):
            price = float(price)
            price_qujian = 0
            if price < 5:
                price_qujian = 6
            elif 5 <= price < 15:
                price_qujian = 5
            elif 15 <= price < 29:
                price_qujian = 0
            elif 29 <= price < 36:
                price_qujian = 1
            elif 36 <= price < 49:
                price_qujian = 2
            elif 49 <= price < 65:
                price_qujian = 3
            elif price >= 65:
                price_qujian = 4
            spu_price_qujian[spu] = price_qujian

    if os.path.exists('./data/FM_data/'+train_file+'.txt'):
        pass
    else:
        with open('./data/FM_data/'+train_file+'.txt', 'w') as f:
            for tup in zip(train['user_id'], train['wm_food_spu_id'], train['ord_period_name'], train['aor_id']):
                f.write(str(1) + ',' + '0:' + str(tup[0]) + ':1,' + '1:' + str(tup[1] + spu_off) + ':1,' + \
                        '2:' + str(tup[2] + period_off) + ':1,' + '3:' + str(tup[3] + aor_off) + ':1,' + \
                        '4:' + str(spu_of_poi_dict[tup[1]][0] + poi_off) + ':1,' + \
                        '5:' + str(user_avg_pay_amt_dict[tup[0]] + avg_pay_amt_off) + ':1,' + \
                        '6:' + str(spu_price_qujian[tup[1]] + spu_price_qujian_off) + ':1')
                f.write('\n')
                while True:
                    neg_item = np.random.randint(low=0, high=spu_num, size=1)[0]
                    if neg_item not in user_spu[tup[0]]:
                        break
                if len(spu_of_poi_dict[neg_item]) == 0:  # 这种情况是这个spu在目前没有商家，因为没有用户购买的记录。这里的处理方法是不加入商家这个特征
                    f.write(str(0) + ',' + '0:' + str(tup[0]) + ':1,' + '1:' + str(neg_item + spu_off) + ':1,' + \
                            '2:' + str(tup[2] + period_off) + ':1,' + '3:' + str(tup[3] + aor_off) + ':1,' + \
                            '5:' + str(user_avg_pay_amt_dict[tup[0]] + avg_pay_amt_off) + ':1,' + \
                            '6:' + str(spu_price_qujian[neg_item] + spu_price_qujian_off) + ':1')
                else:
                    f.write(str(0) + ',' + '0:' + str(tup[0]) + ':1,' + '1:' + str(neg_item + spu_off) + ':1,' + \
                            '2:' + str(tup[2] + period_off) + ':1,' + '3:' + str(tup[3] + aor_off) + ':1,' + \
                            '4:' + str(spu_of_poi_dict[neg_item][0] + poi_off) + ':1,' + \
                            '5:' + str(user_avg_pay_amt_dict[tup[0]] + avg_pay_amt_off) + ':1,' + \
                            '6:' + str(spu_price_qujian[neg_item] + spu_price_qujian_off) + ':1')
                f.write('\n')

    if args.model == 'fm':
        model = xl.create_fm()
    elif args.model == 'lr':
        model = xl.create_linear()
    elif args.model == 'ffm':
        model = xl.create_ffm()

    if args.pre_train == 'pre':
        model.setPreModel('./weights/'+args.model_file)

    param = {'task': 'binary', 'lr': args.lr, 'lambda': args.l2, 'metric': 'acc', 'epoch': args.epoch,
             'stop_window': args.stop_window, 'k': args.dim, 'init': args.init}

    if mode == 'evaluate':
        pass
    else:
        # model.setQuiet() # 设置中间不输出任何评价指标，加快训练速度。
        model.setTrain("./data/FM_data/"+train_file+".txt")
        # model.setValidate("./data/FM_data/test.txt")

        model.fit(param, "./weights/submit-model_" + id + ".out")

    model.setSigmoid()

    user_bprmf_recall_spu = defaultdict(list)
    for line in open('../solution_1/result/'+bprmf_out+'.txt'):
        lines = line.strip().split('\t')
        user_bprmf_recall_spu[int(lines[0])].extend([int(a) for a in lines[1:]])

    unique_test = test[-test.duplicated(['wm_order_id'])]
    if os.path.exists('./data/FM_data/'+bprmf_out+'-'+train_file+'.txt'):
        pass
    else:
        with open('./data/FM_data/'+bprmf_out+'-'+train_file+'.txt', 'w') as f:
            for tup in zip(unique_test['user_id'], unique_test['ord_period_name'], unique_test['aor_id']):
                for spu in user_bprmf_recall_spu[tup[0]]:
                    if len(spu_of_poi_dict[spu]) == 0:
                        f.write(str(1) + ',' + '0:' + str(tup[0]) + ':1,' + '1:' + str(spu + spu_off) + ':1,' + \
                                '2:' + str(tup[1] + period_off) + ':1,' + '3:' + str(tup[2] + aor_off) + ':1,' + \
                            '5:' + str(user_avg_pay_amt_dict[tup[0]] + avg_pay_amt_off) + ':1,'+ \
                            '6:' + str(spu_price_qujian[spu] + spu_price_qujian_off) + ':1')
                    else:
                        f.write(str(1) + ',' + '0:' + str(tup[0]) + ':1,' + '1:' + str(spu + spu_off) + ':1,' + \
                                '2:' + str(tup[1] + period_off) + ':1,' + '3:' + str(tup[2] + aor_off) + ':1,' + \
                                '4:' + str(spu_of_poi_dict[spu][0] + poi_off) + ':1,' + \
                            '5:' + str(user_avg_pay_amt_dict[tup[0]] + avg_pay_amt_off) + ':1,'+ \
                            '6:' + str(spu_price_qujian[spu] + spu_price_qujian_off) + ':1')
                    f.write('\n')

    model.setTest('./data/FM_data/'+bprmf_out+'-'+train_file+'.txt') # 这个有两个点要注意：如果已经存在不重复生成；bin文件生成后是否要删除。
    model.predict("./weights/submit-model_" + id + ".out", "./result/submit-output" + id + ".txt")

    fm_out = pd.read_table("./result/submit-output" + id + ".txt", header=None)
    index = 0
    recall_num = 50
    # 存放按照订单重排的结果
    order_rank = {}

    with open('./submit/submit-'+args.model+'_50_spu_' + id + ".txt", 'w') as f:
        for tup in zip(unique_test['user_id'], unique_test['wm_order_id']):
            spu_score = {}
            for spu in user_bprmf_recall_spu[tup[0]]:
                spu_score[spu] = fm_out[0][index]
                index += 1
            # 控制召回的数目
            count = 0
            spu_score_c = {}
            for key in spu_score:
                spu_score_c[key] = spu_score[key]
                count += 1
                if count == recall_num:
                    break
            S = sorted(spu_score_c.items(), key=lambda item: item[1], reverse=True)
            order_rank[tup[1]] = S
        for tup in test['wm_order_id']:
            f.write(str(tup))

            for s in order_rank[tup]:
                f.write('\t')
                f.write(str(s[0]))
            f.write('\n')
    # 生成提交给网站的结果
    with open('./submit/submit-'+args.model+'_5_spu_' + id + '.txt', 'w') as f:
        for tup in test['wm_order_id']:
            f.write(str(tup))
            count = 0
            for s in order_rank[tup]:
                f.write('\t')
                f.write(str(s[0]))
                count += 1
                if count == 5:
                    break
            f.write('\n')
    print('完成训练')
